% Use a value-function-iteration procedure to obtain the policy function
% under full-employment stabilizing float
%(c) Stephanie Schmitt-Grohe and Martin Uribe, 2010.

clear; clc;
addpath('Auxiliary_functions/')
if ~isfolder('Policy_functions')
    mkdir('.','Policy_functions');
end    

[MC, par, grid]=get_parameterization;

n = grid.n_y*grid.n_d; %number of states

filename = 'Policy_functions/vfi_float_g_low_beta_large_grid';

d_try = repmat(grid.d_grid',grid.n_y,1);
d_try = d_try(:);

r_grid = repmat(grid.r_level,grid.n_d,1);
y_T_grid = repmat(grid.y_level,grid.n_d,1);
g_grid = repmat(grid.g_level,grid.n_d,1);

%construct cTtry = all possibilities for yT-d+dp/(1+r)
c_T_try =y_T_grid-d_try+grid.d_grid'./(1+r_grid);

if min(max(c_T_try,[],2))<0
    error('Natural debt limit violated')
end

c_try = (par.omega*c_T_try.^(1-1/par.xi) + (1-par.omega)*(par.hbar^par.alfa-g_grid).^(1-1/par.xi)).^(1/(1-1/par.xi));%composite consumption

u_try = (c_try.^(1-par.sigg)-1)/(1-par.sigg);
u_try(c_T_try<=0) = -inf;

clear('c_try','cT_try')

v = zeros(grid.n_y,grid.n_d);
v_new = v;
dp_index = zeros(grid.n_y,grid.n_d);
dp_indexnew = dp_index;

dist = 1; 

tic
while dist>1e-4    
    [v_new(:), dp_indexnew(:)] = max(u_try+par.betta*repmat(MC.P_trans*v,grid.n_d,1),[],2);   
    dist = max(abs(v_new(:)-v(:))) + max(abs(dp_indexnew(:)-dp_index(:)))
    v = v_new;    
    %accelerate convergence of v if policy function (dp_index)  has converged
    if (isequal(dp_index,dp_indexnew)) && (dist>1e-8)
        distv = 1;
        picker = sub2ind([n grid.n_d],(1:n)',dp_index(:));
        u = u_try(picker);
        v1 = zeros(grid.n_y,grid.n_d);
        while distv>1e-8            
            temp=repmat(MC.P_trans*v,grid.n_d,1);            
            v1(:) = u + par.betta*temp(picker);            
            distv = max(abs(v(:)-v1(:)))
            v = v1;
        end %while distv        
    end %if dpixold==dpix    
    dp_index = dp_indexnew;    
end %while dist>1e-4
toc
clear('picker','n','u_try','d_try','dp_indexnew','v_new','c_T_try','g_grid','r_grid','temp','v1','u','y_T_grid')

dp = grid.d_grid(dp_index);

save(filename,'-v7.3')